# Copyright (c) 2024-present AI-Labs

from fastapi import APIRouter
from fastapi.responses import FileResponse

import os
import torch
import uuid
import soundfile as sf

from configs import config

from datetime import datetime

from ipex_llm import optimize_model
from transformers import SpeechT5Processor, SpeechT5HifiGan, SpeechT5ForTextToSpeech
from datasets import load_dataset

# 定义路由哦信息
router = APIRouter(
    prefix='/tts/speecht5',
    tags = ['语音生成']
)

# 使用配置文件
model_path = config.service.speecht5.model_path
vocoder_path = config.service.speecht5.vocoder_path
dataset_path = config.service.speecht5.dataset_path

# 加载模型信息
processor = SpeechT5Processor.from_pretrained(model_path)
model = SpeechT5ForTextToSpeech.from_pretrained(model_path)
vocoder = SpeechT5HifiGan.from_pretrained(vocoder_path)

model = optimize_model(model, modules_to_not_convert=["speech_decoder_postnet.feat_out",
                                                      "speech_decoder_postnet.prob_out"]) 
model = model.to('xpu')
vocoder = vocoder.to('xpu')


"""
文本转语音的具体实现过程
"""
def tts_to_file(input):
    # 创建本地路径
    localdir = f"audio/{datetime.now().strftime('%Y-%m-%d')}"
    os.makedirs(f"{config.setting.statics.path}/{localdir}", exist_ok=True)
    localfile = f"{localdir}/{uuid.uuid4()}.wav"

    # 将数据加载到XPU
    inputs = processor(text=input, return_tensors="pt").to('xpu')
    
    # load xvector containing speaker's voice characteristics from a dataset
    embeddings_dataset = load_dataset(dataset_path, split="validation", trust_remote_code=True)
    speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0).to('xpu')

    # 进行模型推理
    with torch.inference_mode():
        # ipex_llm model needs a warmup, then inference time can be accurate
        speech = model.generate_speech(inputs["input_ids"], speaker_embeddings, vocoder=vocoder)

        speech = model.generate_speech(inputs["input_ids"], speaker_embeddings, vocoder=vocoder)
        torch.xpu.synchronize()

    # 模型推理结果保存到本地文件并返回
    sf.write(f"{config.setting.statics.path}/{localfile}", speech.to('cpu').numpy(), samplerate=16000)

    return localfile


"""
对外暴露的接口
"""
@router.post("/audio/speech")
def audio_speech(request: dict):
    # return FileResponse(tts_to_file(request["input"]), media_type="audio/wav")
    return {
                "data": {
                    "url": f"{config.setting.statics.urls}/{tts_to_file(request['input'])}"
                }
           }


"""
对外暴露的接口，返回语音数据的URL地址
"""
@router.post("/generate_audio_url")
def generate_audio_url(request: dict):
    return f"{config.setting.statics.urls}/{tts_to_file(request['input'])}"


"""
对外暴露的接口，返回语音数据文件的二进制数据
"""
@router.post("/generate_audio_file")
def generate_audio_file(request: dict):
    return FileResponse(f"{config.setting.statics.path}/{tts_to_file(request['input'])}", media_type="audio/wav")
